# This class is a wrapper around a numpy matrix that represents the adjacency
# matrix of a graph. It can be any combination of weighted and directed.

# Indices are assumed to be base-0.

# The semantics of the matrix are that rows correspond to nodes and columns
# give the parents (for directed graphs).

# It is designed to work with the Dimacs graph format.

import numpy as np

class DimacsGraph:

    def __init__(self, size, weighted=False, directed=False):
        self.adjacencyMatrix = np.zeros( (size, size) )
        self.weighted = weighted
        self.directed = directed

    def addEdge(self, source, destination, weight=1.0):
        self.adjacencyMatrix[destination, source] = weight
        #print("Adding edge: source '{}', destination '{}', weight '{}'".format(source, destination, weight))

        if not self.directed:
            self.adjacencyMatrix[source, destination] = weight
            #print("Adding edge: source '{}', destination '{}', weight '{}'".format(destination, source, weight))

    def getEdge(self, source, destination):
        return self.adjacencyMatrix[destination, source]

    def getEdgeCount(self):
        return np.count_nonzero(self.adjacencyMatrix)

    def size(self):
        return len(self.adjacencyMatrix)

    def __len__(self):
        return self.size()

    def __getitem__(self, tup):
        x, y = tup
        return self.adjacencyMatrix[x, y]

    # variables is a list of indicators
    # a non-zero value means to keep that variable
    def getSubgrah(self, variables):
        subgraph = DimacsGraph(len(variables), self.weighted, self.directed)
        subgraph.adjacencyMatrix = self.adjacencyMatrix.copy()

        for x in range(len(variables)):
            if variables[x] == 0:
                # clear the row
                subgraph.adjacencyMatix[x,:] = 0

                # clear the column
                subgraph.adjacencyMatrix[:,x] = 0

        return subgraph

    def getReducedSubgraph(self, variables):
        variables = np.nonzero(variables)[0]
        subgraph = DimacsGraph(len(variables), self.weighted, self.directed)
        subgraph.adjacencyMatrix = (self.adjacencyMatrix[variables][:,variables]).copy()
        return subgraph

    def normalize(self):
        ma = np.ma.masked_equal(self.adjacencyMatrix, 0.0, copy=False)
        min = ma.min()
        max = ma.max()
        r = max - min

        for x in range(len(self)):
            for y in range(len(self)):
                if self.adjacencyMatrix[x,y] != 0:
                    # subtract from 1.1 because even the min value is more important than a 0
                    self.adjacencyMatrix[x, y] = 1.1 - (self.adjacencyMatrix[x, y] - min) / r

    def symmetrizeMax(self):
        self.directed = False
        for x in range(len(self)):
            for y in range(x+1, len(self)):
                val = max(self.adjacencyMatrix[x, y], self.adjacencyMatrix[y, x])
                self.adjacencyMatrix[x, y] = val
                self.adjacencyMatrix[y, x] = val

    def scale(self, multiplier = 1):
        if multiplier != 1:
            self.adjacencyMatrix = np.multiply(self.adjacencyMatrix, multiplier)

    
    def floor(self):
        self.adjacencyMatrix = np.floor(self.adjacencyMatrix)
      
    def toUnweighted(self):
        """ This function changes the graph to be unweighted. In particular,
            it changes the dtype of the adjacency matrix to int, and changes
            all of the nonzero weights to int(1).

            N.B. This function returns a copy of the graph.
        """
        unweighted = DimacsGraph(len(self), weighted=False, directed=self.directed)
        m_nonzero = self.adjacencyMatrix != 0
        unweighted.adjacencyMatrix = np.zeros_like(self.adjacencyMatrix, dtype=int)
        unweighted.adjacencyMatrix[m_nonzero] = 1
        return unweighted




